package com.asgard.crypto;

import org.apache.commons.codec.binary.Base64;
import org.apache.commons.lang3.ArrayUtils;
import org.apache.commons.lang3.Validate;
import org.bouncycastle.asn1.sec.SECNamedCurves;
import org.bouncycastle.asn1.x9.X9ECParameters;
import org.bouncycastle.asn1.x9.X9IntegerConverter;
import org.bouncycastle.crypto.params.ECDomainParameters;
import org.bouncycastle.crypto.params.ECPrivateKeyParameters;
import org.bouncycastle.crypto.params.ECPublicKeyParameters;
import org.bouncycastle.crypto.signers.ECDSASigner;
import org.bouncycastle.jcajce.provider.asymmetric.ec.BCECPrivateKey;
import org.bouncycastle.jcajce.provider.asymmetric.ec.BCECPublicKey;
import org.bouncycastle.jcajce.provider.asymmetric.ec.IESCipher;
import org.bouncycastle.jcajce.provider.asymmetric.util.EC5Util;
import org.bouncycastle.jce.ECNamedCurveTable;
import org.bouncycastle.jce.ECPointUtil;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.jce.spec.ECNamedCurveParameterSpec;
import org.bouncycastle.jce.spec.ECPrivateKeySpec;
import org.bouncycastle.jce.spec.ECPublicKeySpec;
import org.bouncycastle.math.ec.ECAlgorithms;
import org.bouncycastle.math.ec.ECCurve;
import org.bouncycastle.math.ec.ECCurve.Fp;
import org.bouncycastle.math.ec.ECPoint;
import org.bouncycastle.util.Arrays;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import java.io.ByteArrayOutputStream;
import java.math.BigInteger;
import java.security.*;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

import static com.asgard.crypto.CryptoUtils.encodeRipeMD160;
import static com.asgard.crypto.CryptoUtils.encodeSha256;

/**
 * 该类是椭圆曲线加密的辅助类。
 */
public class Elliptic {

    private static ECNamedCurveParameterSpec ecGenSpec;
    private static KeyPairGenerator g;
    private static KeyFactory fact;
    private static X9ECParameters x9;
    private static BouncyCastleProvider bc = null;

    static {
        try {
            /*
             * Elliptic 类被多处使用，其一是云钱包中用户注册时生成 ECKey； 其二是公共账簿中用于生成 Address 对应的
			 * Hash 160.
			 * 
			 * 在两个系统， 云钱包和公共帐薄，分别以war的形式部署在同一个tomcat实例时， 使用Elliptic类会抛出
			 * NoClassDefFoundError. 这是一个J2SE实现问题。
			 * 
			 * 以下是 Security.addProvider() 执行逻辑， First, if there is a security
			 * manager, its checkSecurityAccess method is called with the string
			 * "insertProvider."+provider.getName() to see if it's ok to add a
			 * new provider. If the default implementation of
			 * checkSecurityAccess is used (i.e., that method is not overriden),
			 * then this will result in a call to the security manager's
			 * checkPermission method with a
			 * SecurityPermission("insertProvider."+provider.getName())
			 * permission.
			 * 
			 * 所以，NoClassDefFoundError 错误的原因在于，再次添加provider(相对于 JVM)时， throw
			 * SecurityException if a security manager exists and its
			 * SecurityManager.checkSecurityAccess(java.lang.String) method
			 * denies access to add a new provider.
			 * 
			 * The recommended way to avoid that trouble is to have bouncy
			 * castle classes in your JVM ClassPath or in your container
			 * ClassPath. You have to remove it from your application. Now you
			 * need to register BC provider with an alternate option to the
			 * static initializer. WebLogic provides ways to trigger code at
			 * server startup (I have used server startup class), this code will
			 * be responsible to register JSSE providers for the whole
			 * server/JVM lifetime.
			 * 
			 * An alternate option is to add the following line in JRE
			 * java.security file with bouncy castle jar in jre/lib/ext but I do
			 * not like that way because it may be lost when updating:
			 * security.provider.7=org.bouncycastle.jce.provider.
			 * BouncyCastleProvider
			 * 
			 * {reference link:
			 * http://stackoverflow.com/questions/10379799/bouncycastle-cannot-
			 * load-when-i-redeploy-application}
			 */
            if (bc == null) {
                bc = new BouncyCastleProvider();
                Security.removeProvider(bc.getName());
                Security.addProvider(bc);
                ecGenSpec = ECNamedCurveTable.getParameterSpec("prime256v1");
                x9 = SECNamedCurves.getByName("prime256v1");
                g = KeyPairGenerator.getInstance("ECDSA", "BC");
                g.initialize(ecGenSpec, new SecureRandom());
                fact = KeyFactory.getInstance("ECDSA", "BC");
            }
        } catch (Exception e) {
            throw new RuntimeException("无法初始化", e);
        }
    }

    public static void init() {
        // do nothing
    }

    /**
     * 生成新的椭圆曲线加密密钥对。
     *
     * @return 生成的新密钥对
     */
    public static KeyPair generateKeyPair() {
        KeyPair pair = g.generateKeyPair();
        return pair;
    }

    /**
     * 用指定的私钥对消息进行签名。
     *
     * @param privKey 私钥
     * @param data    数据
     * @return 签名
     */
    public static byte[] sign(PrivateKey privKey, byte[] data) {
        try {
            Signature ecdsaSign = Signature.getInstance("SHA256withECDSA", "BC");
            ecdsaSign.initSign(privKey);
            ecdsaSign.update(data);
            byte[] signature = ecdsaSign.sign();
            return signature;
        } catch (InvalidKeyException e) {
            throw new RuntimeException(e);
        } catch (SignatureException e) {
            throw new RuntimeException(e);
        } catch (NoSuchAlgorithmException | NoSuchProviderException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 检验签名是否正确。
     *
     * @param pubKey    公钥
     * @param data      数据
     * @param signature 签名
     * @return 签名是否正确
     */
    public static boolean verify(PublicKey pubKey, byte[] data, byte[] signature) {
        try {
            Signature ecdsaVerify = Signature.getInstance("SHA256withECDSA", "BC");
            ecdsaVerify.initVerify(pubKey);
            ecdsaVerify.update(data);
            boolean result = ecdsaVerify.verify(signature);
            return result;
        } catch (NoSuchAlgorithmException | NoSuchProviderException e) {
            throw new RuntimeException(e);
        } catch (InvalidKeyException e) {
            throw new RuntimeException(e);
        } catch (SignatureException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 按照PKCS8格式将私钥转换成为字节数组。
     *
     * @param privKey 私钥
     * @return 字节数组
     */
    public static byte[] getBytesPKCS8(PrivateKey privKey) {
        BCECPrivateKey priKey = (BCECPrivateKey) privKey;
        return priKey.getEncoded();
    }

    /**
     * 按照一般格式将私钥转换为字节数组。
     *
     * @param privKey 私钥
     * @return 字节数组
     */
    public static byte[] getBytesNormal(PrivateKey privKey) {
        BCECPrivateKey priKey = (BCECPrivateKey) privKey;
        BigInteger s = priKey.getS();
        byte[] b = s.toByteArray();
        return toByte32(b);
    }

    private static byte[] toByte32(byte[] b) {
        if (b.length == 32)
            return b;
        else if (b.length > 32) {
            Validate.validState(b[0] == 0);
            return ArrayUtils.subarray(b, b.length - 32, b.length);
        } else {
            byte[] b2 = new byte[32];
            System.arraycopy(b, 0, b2, 32 - b.length, b.length);
            return b2;
        }
    }

    /**
     * 按照X509格式将公钥转换为字节数组。
     *
     * @param pubKey 公钥
     * @return 字节数组
     */
    public static byte[] getBytesX509(PublicKey pubKey) {
        BCECPublicKey pubKey2 = (BCECPublicKey) pubKey;
        return pubKey2.getEncoded();
    }

    /**
     * 按照一般格式将公钥转换成为字节数组。
     *
     * @param pubKey 公钥
     * @return 字节数组
     */
    public static byte[] getBytesNormal(PublicKey pubKey) {
        BCECPublicKey pubKey2 = (BCECPublicKey) pubKey;
        ECPoint p = EC5Util.convertPoint(pubKey2.getParams(), pubKey2.getW(), false);
        return p.getEncoded(false);
    }

    /**
     * 按照压缩格式将公钥转换为字节数组。
     *
     * @param pubKey 公钥
     * @return 字节数组
     */
    public static byte[] getBytesCompact(PublicKey pubKey) {
        BCECPublicKey pubKey2 = (BCECPublicKey) pubKey;
        ECPoint p = EC5Util.convertPoint(pubKey2.getParams(), pubKey2.getW(), false);
        return p.getEncoded(true);
    }

    /**
     * 将字节数组还原为公钥。
     *
     * @param encoded 字节数组
     * @return 公钥
     * @throws InvalidKeySpecException 字节数组格式不正确
     */
    public static PublicKey getPublicKeyAuto(byte[] encoded) throws InvalidKeySpecException {
        if (encoded.length == 65)
            return getPublicKeyNormal(encoded);
        if (encoded.length == 33)
            return getPublicKeyCompact(encoded);
        else
            return getPublicKeyX509(encoded);
    }

    /**
     * 将X509格式的字节数组转换为公钥。
     *
     * @param encoded 字节数组
     * @return 公钥
     * @throws InvalidKeySpecException 字节数组格式不正确
     */
    public static PublicKey getPublicKeyX509(byte[] encoded) throws InvalidKeySpecException {
        PublicKey pubKey = fact.generatePublic(new X509EncodedKeySpec(encoded));
        return pubKey;
    }

    /**
     * 将一般格式的字节数组转换成为公钥。
     *
     * @param encoded 字节数组
     * @return 公钥
     * @throws InvalidKeySpecException 字节数组格式不正确
     */
    public static PublicKey getPublicKeyNormal(byte[] encoded) throws InvalidKeySpecException {
        ECCurve curve = ecGenSpec.getCurve();
        java.security.spec.EllipticCurve ellipticCurve = EC5Util.convertCurve(curve, ecGenSpec.getSeed());
        java.security.spec.ECPoint point = ECPointUtil.decodePoint(ellipticCurve, encoded);
        java.security.spec.ECParameterSpec params2 = EC5Util.convertSpec(ellipticCurve, ecGenSpec);
        java.security.spec.ECPublicKeySpec keySpec = new java.security.spec.ECPublicKeySpec(point, params2);
        return fact.generatePublic(keySpec);
    }

    /**
     * 将压缩格式的字节数组转换为公钥。
     *
     * @param encoded 字节数组
     * @return 公钥
     * @throws InvalidKeySpecException 字节数组格式不正确
     */
    public static PublicKey getPublicKeyCompact(byte[] encoded) throws InvalidKeySpecException {
        ECCurve curve = ecGenSpec.getCurve();
        ECPoint q = curve.decodePoint(encoded);
        return fact.generatePublic(new ECPublicKeySpec(q, ecGenSpec));
    }

    /**
     * 将字节数组还原为私钥。
     *
     * @param encoded 字节数组
     * @return 私钥
     * @throws InvalidKeySpecException 字节数组格式不正确
     */
    public static PrivateKey getPrivateKeyAuto(byte[] encoded) throws InvalidKeySpecException {
        if (encoded.length == 32)
            return getPrivateKeyNormal(encoded);
        else
            return getPrivateKeyPKCS8(encoded);
    }

    /**
     * 将PKCS8格式的字节数组还原为私钥。
     *
     * @param encoded 字节数组
     * @return 私钥
     * @throws InvalidKeySpecException 字节数组格式不正确
     */
    public static PrivateKey getPrivateKeyPKCS8(byte[] encoded) throws InvalidKeySpecException {
        PrivateKey privKey = fact.generatePrivate(new PKCS8EncodedKeySpec(encoded));
        return privKey;
    }

    /**
     * 将一般格式的字节数组还原为私钥
     *
     * @param encoded 字节数组
     * @return 私钥
     * @throws InvalidKeySpecException 字节数组格式不正确
     */
    public static PrivateKey getPrivateKeyNormal(byte[] encoded) throws InvalidKeySpecException {
        BigInteger d = new BigInteger(1, encoded);
        return fact.generatePrivate(new ECPrivateKeySpec(d, ecGenSpec));
    }

    /**
     * 将一般格式的字节数组表示的公钥编码成地址。
     *
     * @param pubKey 公钥
     * @return 地址
     */
    @Deprecated
    public static String publicKeyToAddressNormal(PublicKey pubKey) {
        return publicKeyToAddress(Elliptic.getBytesNormal(pubKey));
    }

    /**
     * 将压缩格式的字节数组表示的公钥编码成地址。
     *
     * @param pubKey 公钥
     * @return 地址
     */
    public static String publicKeyToAddressCompact(PublicKey pubKey) {
        return publicKeyToAddress(Elliptic.getBytesCompact(pubKey));
    }

    /**
     * 将公钥转换成地址。
     *
     * @param pubKey 字节数据表示的公钥
     * @return 地址
     */
    public static String publicKeyToAddress(byte[] pubKey) {
        byte[] hash = encodeRipeMD160(encodeSha256(pubKey));
        return publicKeyHashToAddress(hash);
    }

    /**
     * 将公钥哈希转换成地址。
     *
     * @param pubKeyHash 公钥哈希
     * @return 地址
     */
    public static String publicKeyHashToAddress(byte[] pubKeyHash) {
        byte[] d = new byte[pubKeyHash.length + 1];
        System.arraycopy(pubKeyHash, 0, d, 1, pubKeyHash.length);
        String url = CryptoUtils.encodeBase58Verified(d);
        return url;
    }

    /**
     * 将地址转换成公钥哈希。
     *
     * @param address 地址
     * @return 公钥哈希
     */
    public static byte[] addressToPublicKeyHash(String address) {
        byte[] d = CryptoUtils.decodeBase58Verified(address);
        return ArrayUtils.subarray(d, 1, d.length);
    }

    /**
     * 对文本消息的签名进行校验。
     *
     * @param message   文本消息
     * @param base64Sig 对文本消息的签名
     * @param address   地址
     * @return 交易结果
     */
    public static boolean verifySignedMessage(String message, String base64Sig, String address) {
        try {
            byte[] sigBytes = Base64.decodeBase64(base64Sig);
            String reEncode = Base64.encodeBase64String(sigBytes);
            if (!reEncode.equals(base64Sig)) {
                return false;
            }
            byte[] msgBytes = prepareMessage(message);
            int first = (sigBytes[0] - 27);
            boolean fCompressed = (first & 4) != 0;
            int rec = first & 3;
            BigInteger[] sig = parseSig(sigBytes);
            byte[] msgHash = CryptoUtils.encodeSha256Twice(msgBytes);
            ECPoint Q = recover(msgHash, sig, rec, true);
            byte[] qEnc = Q.getEncoded(fCompressed);
            String qAddr = publicKeyToAddress(qEnc);
            if (!qAddr.equals(address)) {
                return false;
            }
            return verifySignature(Q, msgHash, sig);
        } catch (Exception e) {
            return false;
        }
    }

    private static byte[] prepareMessage(String msg) {
        return msg.getBytes();
    }

    private static BigInteger[] parseSig(byte[] sigBytes) {
        byte[] tmp = ArrayUtils.subarray(sigBytes, 1, 33);
        BigInteger r = new BigInteger(1, tmp);
        tmp = ArrayUtils.subarray(sigBytes, 33, 65);
        BigInteger s = new BigInteger(1, tmp);
        return new BigInteger[]{r, s};
    }

    private static ECPoint recover(byte[] msgHash, BigInteger[] sig, int recid, boolean check) {
        // 曲线基本信息
        BigInteger r = sig[0], s = sig[1];
        Fp curve = (Fp) x9.getCurve();
        BigInteger order = x9.getN();
        BigInteger x = r;
        if ((recid & 2) != 0)
            x = x.add(order);
        if (x.compareTo(curve.getQ()) >= 0)
            throw new IllegalArgumentException();
        X9IntegerConverter x9Converter = new X9IntegerConverter();
        byte[] xEnc = x9Converter.integerToBytes(x, x9Converter.getByteLength(curve));
        byte[] compEncoding = new byte[xEnc.length + 1];
        compEncoding[0] = (byte) (0x02 + (recid & 1));
        System.arraycopy(xEnc, 0, compEncoding, 1, xEnc.length);
        ECPoint R = x9.getCurve().decodePoint(compEncoding);
        if (check) {
            ECPoint O = R.multiply(order);
            if (!O.isInfinity())
                throw new IllegalArgumentException();
        }
        BigInteger e = calculateE(order, msgHash);
        BigInteger rInv = r.modInverse(order);
        BigInteger srInv = s.multiply(rInv).mod(order);
        BigInteger erInv = e.multiply(rInv).mod(order);
        return ECAlgorithms.sumOfTwoMultiplies(R, srInv, x9.getG().negate(), erInv);
    }

    private static BigInteger calculateE(BigInteger n, byte[] message) {
        int messageBitLength = message.length * 8;
        byte[] tmp = Arrays.clone(message);
        BigInteger trunc = new BigInteger(1, tmp);
        if (n.bitLength() < messageBitLength)
            trunc = trunc.shiftRight(messageBitLength - n.bitLength());
        return trunc;
    }

    private static boolean verifySignature(ECPoint Q, byte[] hash, BigInteger[] sig) throws InvalidKeySpecException {
        ECDomainParameters ec = new ECDomainParameters(x9.getCurve(), x9.getG(), x9.getN(), x9.getH(), x9.getSeed());
        ECPublicKeyParameters publicKey = new ECPublicKeyParameters(Q, ec);
        boolean succ = verifySignature(publicKey, hash, sig);
        System.out.println(succ ? "Verify succeed" : "Verify fail: Signature doesn't match");
        return succ;
    }

    private static boolean verifySignature(ECPublicKeyParameters publicKey, byte[] hash, BigInteger[] sig) {
        ECDSASigner signer = new ECDSASigner();
        signer.init(false, publicKey);
        return signer.verifySignature(hash, sig[0], sig[1]);
    }

    /**
     * 利用文本消息专用的压缩格式对文本消息进行签名。
     *
     * @param keyPair      密钥对
     * @param data         文本消息
     * @param isCompressed 是否压缩
     * @return 签名
     */
    public static byte[] signCompact(KeyPair keyPair, byte[] data, boolean isCompressed) {
        BCECPrivateKey priKey = (BCECPrivateKey) keyPair.getPrivate();
        BCECPublicKey pubKey = (BCECPublicKey) keyPair.getPublic();
        ECDSASigner signer = new ECDSASigner();
        BigInteger d = priKey.getD();
        ECDomainParameters ec = new ECDomainParameters(x9.getCurve(), x9.getG(), x9.getN(), x9.getH(), x9.getSeed());
        signer.init(true, new ECPrivateKeyParameters(d, ec));
        BigInteger[] rs = signer.generateSignature(data);
        try {
            ByteArrayOutputStream bout = new ByteArrayOutputStream();
            int first = 27;
            if (isCompressed)
                first += 4;
            int rec = getRecId(data, rs, pubKey, isCompressed);
            first += rec;
            bout.write(first);
            bout.write(toByte32(rs[0].toByteArray()));
            bout.write(toByte32(rs[1].toByteArray()));
            return bout.toByteArray();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    private static int getRecId(byte[] msgHash, BigInteger[] sig, BCECPublicKey pubKey, boolean isCompressed) {
        ECPoint Q0 = pubKey.getQ();
        for (int rec = 0; rec < 4; rec++) {
            try {
                ECPoint Q = recover(msgHash, sig, rec, true);
                if (Q.equals(Q0))
                    return rec;
            } catch (Exception e) {
            }
        }
        return 0;
    }

    /**
     * 对文本消息进行签名。
     *
     * @param keyPair      密钥对
     * @param message      文本消息
     * @param isCompressed 是否压缩
     * @return 对文本消息生成的签名
     */
    public static String signMessage(KeyPair keyPair, String message, boolean isCompressed) {
        byte[] data = prepareMessage(message);
        byte[] hash = CryptoUtils.encodeSha256Twice(data);
        byte[] sig = signCompact(keyPair, hash, isCompressed);
        return Base64.encodeBase64String(sig);
    }

    /**
     * 根据私钥生成公钥。
     *
     * @param privKey 私钥
     * @return 公钥
     * @throws InvalidKeySpecException 密钥格式不正确
     */
    public static PublicKey generatePublicKey(PrivateKey privKey) throws InvalidKeySpecException {
        BCECPrivateKey priKey = (BCECPrivateKey) privKey;
        BigInteger d = priKey.getD();
        ECPoint G = ecGenSpec.getG();
        ECPoint Q = G.multiply(d);
        return fact.generatePublic(new ECPublicKeySpec(Q, ecGenSpec));
    }

    /**
     * 用ECC公钥加密数据。
     *
     * @param data      数据
     * @param publicKey 公钥
     * @return 加密后的数据
     */
    public static byte[] encrypt(byte[] data, PublicKey publicKey) {
        try {
            IESCipher.ECIES c1 = new IESCipher.ECIES();
            c1.engineInit(Cipher.ENCRYPT_MODE, publicKey, new SecureRandom());
            return c1.engineDoFinal(data, 0, data.length);
        } catch (InvalidKeyException e) {
            throw new RuntimeException(e);
        } catch (BadPaddingException e) {
            throw new RuntimeException(e);
        } catch (IllegalBlockSizeException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 用ECC私钥解密数据。
     *
     * @param encryptedData 加密后的数据
     * @param privateKey    私钥
     * @return 解密后的数据
     */
    public static byte[] decrypt(byte[] encryptedData, PrivateKey privateKey) {
        try {
            IESCipher.ECIES c1 = new IESCipher.ECIES();
            c1.engineInit(Cipher.DECRYPT_MODE, privateKey, new SecureRandom());
            return c1.engineDoFinal(encryptedData, 0, encryptedData.length);
        } catch (InvalidKeyException e) {
            throw new RuntimeException(e);
        } catch (BadPaddingException e) {
            throw new RuntimeException(e);
        } catch (IllegalBlockSizeException e) {
            throw new RuntimeException(e);
        }
    }

}
